import numpy as np

# 模拟向量数据
np.random.seed(42)
data_vectors = np.random.rand(100, 5)  # 100个5维向量


# -------------------------------
# 分区技术演示
# -------------------------------
class Partitioning:
    def __init__(self, data_vectors, num_partitions):
        self.data_vectors = data_vectors
        self.num_partitions = num_partitions
        self.partitions = {i: [] for i in range(num_partitions)}

    def assign_to_partitions(self):
        """
        根据第一个维度的值将向量分区
        """
        for i, vector in enumerate(self.data_vectors):
            partition_id = int(vector[0] * self.num_partitions)
            self.partitions[partition_id].append(i)

    def search_in_partition(self, partition_id, query_vector, top_k=5):
        """
        在指定分区中检索最相似的向量
        """
        if partition_id not in self.partitions or not self.partitions[partition_id]:
            return []
        partition_vectors = [self.data_vectors[i] for i in self.partitions[partition_id]]
        distances = [np.linalg.norm(query_vector - vec) for vec in partition_vectors]
        nearest_indices = np.argsort(distances)[:top_k]
        return [self.partitions[partition_id][i] for i in nearest_indices]


# 创建分区
num_partitions = 10
partitioning = Partitioning(data_vectors, num_partitions)
partitioning.assign_to_partitions()

# 在分区中检索
partition_to_search = 3
query_vector = np.array([0.5, 0.5, 0.5, 0.5, 0.5])  # 查询向量
nearest_in_partition = partitioning.search_in_partition(partition_to_search, query_vector, top_k=3)
print(f"分区 {partition_to_search} 中最相似的向量ID:", nearest_in_partition)

"""
在向量检索中，除了基于单一维度的简单分区，常见的分区方式还有以下几类，适用于不同场景和数据特性：

1. **空间划分（如网格分区）**  
   将整个向量空间划分为多个互不重叠的网格单元，每个向量根据其在高维空间中的坐标落入对应的网格。查询时只需搜索查询向量所在网格及相邻网格，减少计算量。

2. **聚类分区（如K-Means聚类）**  
   通过聚类算法（如K-Means）将相似向量聚为多个簇（分区），每个簇作为一个分区。查询时先找到与查询向量最相似的簇，再在该簇内精细检索，适合向量分布有明显聚类特性的场景。

3. **分层分区（如树结构分区）**  
   采用树结构（如KD树、R树）递归划分空间，每层根据某个维度的阈值将当前空间分为子空间，形成多层级分区。查询时沿树路径快速定位可能包含相似向量的分区，适合高维数据但维度不极端（如<20维）的场景。

4. **随机分区（如哈希分区）**  
   通过哈希函数（如局部敏感哈希LSH）将向量映射到不同桶（分区），相似向量有更高概率落入同一桶。查询时只需检索查询向量映射到的桶，适合大规模数据的近似检索。

5. **基于核心维度组合的分区**  
   选取多个重要维度（而非单一维度），通过组合这些维度的范围划分分区，比单一维度分区更能反映向量的整体分布。

这些方式的核心目的是通过“粗筛+精查”减少检索时的计算量，平衡效率与精度，具体选择需结合数据维度、分布特点及检索需求。
"""
